package io.github.hadyang.leetcode.offer;

import org.junit.Test;

/** @author haoyang.shi */
public class CountOfSortedArray {

  @Test
  public void test() {
    System.out.println(countOfSortedArray(new int[] {1, 2, 3, 3, 3, 3, 4, 4, 5}, 3));
    System.out.println(countOfSortedArray(new int[] {1, 2, 3, 3, 3, 3, 4, 4, 5}, 4));
    System.out.println(countOfSortedArray(new int[] {1, 2, 3, 3, 3, 3, 4, 4, 5}, 1));

    System.out.println(countOfSortedArray2(new int[] {1, 2, 3, 3, 3, 3, 4, 4, 5}, 3));
    System.out.println(countOfSortedArray2(new int[] {1, 2, 3, 3, 3, 3, 4, 4, 5}, 4));
    System.out.println(countOfSortedArray2(new int[] {1, 2, 3, 3, 3, 3, 4, 4, 5}, 1));
  }

  public int countOfSortedArray(int[] nums, int n) {
    if (nums == null || nums.length == 0) return 0;

    int s = 0, e = nums.length - 1;
    while (s < e) {
      if (nums[s] != n) s++;

      if (nums[e] != n) e--;

      if (nums[s] == nums[e]) break;
    }

    if (nums[s] != nums[e]) return 0;

    return e - s + 1;
  }

  public int countOfSortedArray2(int[] nums, int n) {
    if (nums == null || nums.length == 0) return 0;

    int firstN = getFirstN(nums, n);
    int lastN = getLastN(nums, n);

    return lastN - firstN + 1;
  }

  private int getFirstN(int[] nums, int n) {
    int s = 0, e = nums.length - 1;

    int mid = -1;
    while (s <= e) {
      mid = (s + e) / 2;

      if (mid > 0 && nums[mid - 1] == n) {
        e = mid - 1;
        continue;
      }

      if (nums[mid] > n) {
        e = mid - 1;
        continue;
      }

      if (nums[mid] < n) {
        s = mid + 1;
        continue;
      }

      break;
    }

    return mid;
  }

  private int getLastN(int[] nums, int n) {
    int s = 0, e = nums.length - 1;

    int mid = -1;
    while (s <= e) {
      mid = (s + e) / 2;

      if (mid < nums.length - 1 && nums[mid + 1] == n) {
        s = mid + 1;
        continue;
      }

      if (nums[mid] > n) {
        e = mid - 1;
        continue;
      }

      if (nums[mid] < n) {
        s = mid + 1;
        continue;
      }

      break;
    }

    return mid;
  }
}
